package com.wj.graph;

import com.wj.linear.Queue;
import com.wj.tree.IndexMinPriorityQueue;

import java.util.Arrays;

/**
 * @author wen.jie
 * @date 2021/8/30 11:00
 * prim算法
 */
public class PrimMST {

    //索引代表顶点，值表示当前顶点和最小生成树之间的最短边
    private Edge[] edgeTo;

    //索引代表顶点，值表示当前顶点和最小生成树之间的最短边的权重
    private double[] distTo;

    //索引代表顶点，如果当前顶点已经在树中，则值为true，否则为 false
    private boolean[] marked;

    //存放树中顶点与非树中顶点之间的有效横切边
    private IndexMinPriorityQueue<Double> pq;

    public PrimMST(EdgeWeightedGraph G) {
        this.edgeTo = new Edge[G.V()];
        this.distTo = new double[G.V()];
        Arrays.fill(distTo, Double.POSITIVE_INFINITY);
        this.marked = new boolean[G.V()];
        pq = new IndexMinPriorityQueue<>(G.V());

        //默认让顶点0进入到树中
        distTo[0] = 0.0;
        pq.insert(0, 0.0);
        //遍历索引最小优先队列,拿到最小边和N切边对应的顶点,把该顶点加入到最小生成树中
        while (!pq.isEmpty()) {
            visit(G, pq.delMin());
        }
    }

    //将顶点v添加到最小生成树中，并且更新数据
    private void visit(EdgeWeightedGraph G, int v) {
        marked[v] = true;
        for (Edge e : G.adj(v)) {
            int w = e.other(v);
            if (marked[w])
                continue;
            if (e.weight() < distTo[w]) {
                edgeTo[w] = e;
                distTo[w] = e.weight();
                if (pq.contains(w)) {
                    pq.changeItem(w, e.weight());
                }else {
                    pq.insert(w, e.weight());
                }
            }
        }
    }

    //获取最小生成树的所有边
    public Queue<Edge> edges() {
        Queue<Edge> allEdges = new Queue<>();
        for (Edge edge : edgeTo) {
            if (edge != null) {
                allEdges.enqueue(edge);
            }
        }
        return allEdges;
    }

}
